package cn.bugstack.http;

import cn.bugstack.IAiService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.Map;

/**
 * @Classname OllamaController
 * @Description TODO
 * @Date 2025/10/16 10:59
 * @Author 勾新杰
 */
@RestController
@CrossOrigin(origins = "*")
@RequestMapping("/api/v1/ollama/")
public class OllamaController implements IAiService {

    @Resource
    private OllamaChatClient ollamaChatClient;

    @Resource
    private PgVectorStore pgVectorStore;

    /**
     * 非流式返回
     * @param model 模型
     * @param message 消息
     * @return 响应
     */
    @Override
    @RequestMapping(value = "/generate", method = RequestMethod.GET)
    public ChatResponse generate(@RequestParam String model, @RequestParam String message) {
        return ollamaChatClient.call(new Prompt(message, OllamaOptions.create().withModel(model)));
    }

    /**
     * 流式返回
     * @param model 模型
     * @param message 消息
     * @return 响应
     */
    @Override
    @RequestMapping(value = "/generate-stream", method = RequestMethod.GET)
    public Flux<ChatResponse> generateStream(String model, String message) {
        ArrayList<Message> messages = new ArrayList<>();
        Message systemMessage = new SystemPromptTemplate("""
                                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                                If unsure, simply state that you don't know.
                                Another thing you need to note is that your reply must be in Chinese!
                                DOCUMENTS: {documents}
                """).createMessage(Map.of("documents", pgVectorStore.similaritySearch(SearchRequest.query(message).withTopK(5))));
        messages.add(new UserMessage(message));
        messages.add(systemMessage);
        return ollamaChatClient.stream(new Prompt(messages, OllamaOptions.create().withModel(model)));
    }
}
